package com.apex.flink.source;

import com.apex.flink.FlinkEnvironment;
import com.apex.flink.stream.FlinkStreamSource;
import com.apex.flink.utils.ConfigKeyName;
import com.typesafe.config.Config;
import org.apache.flink.connectors.kudu.batch.KuduRowInputFormat;
import org.apache.flink.connectors.kudu.connector.KuduTableInfo;
import org.apache.flink.connectors.kudu.connector.reader.KuduReaderConfig;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.types.Row;

public class KuduSource implements FlinkStreamSource<Row> {
    private Config config;
    private String tableName;
    private String KUDU_MASTERS;
    @Override
    public DataStream<Row> getData(FlinkEnvironment env) {
        StreamExecutionEnvironment environment = env.getStreamExecutionEnvironment();
        KuduTableInfo tableInfo = KuduTableInfo.forTable(tableName);
        KuduReaderConfig readerConfig = KuduReaderConfig.Builder.setMasters(KUDU_MASTERS).build();
        KuduRowInputFormat inputFormat = new KuduRowInputFormat(readerConfig, tableInfo);
        return environment.createInput(inputFormat, env.getRowTypeInfo());
    }

    @Override
    public void prepare(FlinkEnvironment plugin) {
        if (config.hasPath(ConfigKeyName.FLINK_SOURCE_KUDU_MASTER)){
            KUDU_MASTERS = config.getString(ConfigKeyName.FLINK_SOURCE_KUDU_MASTER);
            tableName = config.getString(ConfigKeyName.FLINK_SOURCE_KUDU_TABLE_NAME);
        }
    }

    @Override
    public Config getConfig() {
        return config;
    }

    @Override
    public void setConfig(Config config) {
        this.config = config;
    }
}
